package com.ruoyi.ade.llm.agent;

import com.alibaba.fastjson.JSONObject;
import com.ruoyi.ade.llm.model.*;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpEntity;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;

public class OllamaAgent extends LLMAgent {

    public static final String LLM_MODEL_QWEN2_7B = "qwen2:7b";
    public static final String LLM_MODEL_QWEN25_7B = "qwen2.5:7b";
    public static final String LLM_MODEL_QWEN25_14B = "qwen2.5:14b";
    public static final String LLM_MODEL_LLAMA31_8B = "llama3.1:8b";

    public static final String url = "http://192.168.254.201:27004";
    public static final String url_chat = url + "/api/chat";
    public static final String url_embed = url + "/api/embed";
    public static final String url_completion = url + "/api/generate";

    private OllamaAgent() {
    }

    // 静态内部类负责实例化，类加载时是延迟的
    private static class SingletonHelper {
        private static final OllamaAgent INSTANCE = new OllamaAgent();
    }

    // 提供全局访问点
    public static OllamaAgent getInstance() {
        return SingletonHelper.INSTANCE;
    }

    /**
     * @param request
     * @return
     */
    public ChatResponse chat(ChatRequest request) {
        String json = post(url_chat, JSONObject.toJSONString(request));
        if (StringUtils.isNoneEmpty(json))
            return JSONObject.parseObject(json, ChatResponse.class);
        return null;
    }

    /**
     * @param request
     * @return
     */
    public CompletionResponse complet(CompletionRequest request) {
        String json = post(url_completion, JSONObject.toJSONString(request));
        if (StringUtils.isNoneEmpty(json))
            return JSONObject.parseObject(json, CompletionResponse.class);
        return null;
    }

    /**
     * 嵌入请求
     *
     * @param request
     * @return
     */
    public EmbeddingResponse embed(EmbeddingRequest request) {
        String json = post(url_embed, JSONObject.toJSONString(request));
        if (StringUtils.isNoneEmpty(json))
            return JSONObject.parseObject(json, EmbeddingResponse.class);
        return null;
    }

    private String post(String url, String jsonRequest) {
        try {
            System.out.println(url);
            System.out.println(jsonRequest);
            CloseableHttpClient httpClient = HttpClients.createDefault();
            HttpPost postRequest = new HttpPost(url);
            StringEntity entity = new StringEntity(jsonRequest, "UTF-8");
            entity.setContentType("application/json");
            postRequest.setEntity(entity);
            CloseableHttpResponse resp = httpClient.execute(postRequest);
            System.out.println(resp.getStatusLine().getStatusCode());
            System.out.println(JSONObject.toJSONString(resp));
            HttpEntity e = resp.getEntity();
            if (e != null && resp.getStatusLine().getStatusCode() == 200) {
                String content = EntityUtils.toString(e);
                return content;
            }
            return null;
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

}
